# 设置两个TDMA节点，6个SA节点
# 距离随机选择1-5时隙，最后就运行10个时隙
import numpy as np
import random

# 计算TDMA节点与DRL节点的传播延迟差
#delay_tag = TDMA_delay - DRL_delay
# 设置DRL节点动作历史长度
DRL_action_length = 30
# 初始化TDMA计数器
TDMA_counter = 0

class ENVIRONMENT(object):
    """
    环境类，用于模拟TDMA和DRL和SA节点的交互
    """

    def __init__(self, state_size=10):
        # 初始化环境状态大小
        self.state_size = state_size
        # 定义动作空间，'w'代表等待，'t'代表传输
        self.action_space = ['w', 't']
        # 获取动作空间的大小
        self.n_actions = len(self.action_space)
        # 设置网络中的节点数量
        self.n_nodes = 3
        # 设置TDMA节点的传播延迟，随机生成
        self.TDMA_delay = random.randint(1, 5)
        # 设置SA节点的传播延迟，随机生成
        # self.SA1_delay = random.randint(1, 5)
        # self.SA2_delay = random.randint(1, 5)
        # 设置DRL节点的传播延迟
        self.DRL_delay = random.randint(1, 5)
        # 设置延迟标签，两个TDMA节点之间与DRL节点之间的延迟差
        #self.delay_tag = delay_tag
        self.delay_tag = self.TDMA_delay - self.DRL_delay

        # 预设TDMA节点的动作序列
        self.TDMAaction = [0, 1, 0, 0, 1, 0, 0, 1, 0, 0]  # case 1
        # 预设SA节点的发送概率
        self.SAaction = 0.1# case 1
        # self.SAaction = [0.8]  # case 1
        #self.SAaction = [0.3, 0.7]  # SA节点比较少的话最优吞吐量比较大
        # 初始化DRL节点的奖励列表
        self.agent_reward_list = []
        # 初始化sa节点的奖励列表
        self.sa_reward_list = []
        # 初始化TDMA节点的奖励列表
        self.tdma_reward_list = []
        # 初始化总奖励列表
        self.reward_list = []
        # 初始化DRL节点的观察列表
        self.DRL_observation_list = []


        # 初始化观察奖励计数器
        self.observation_reward_counter = 0

        # 对于Slotted ALOHA节点 计算z
        # 当只有一个SA结点的时候，计算z的公式简化成了 1-q-q

        self.z = 1-2*self.SAaction

        print("z=",self.z)
    # 重置环境状态
    def reset(self):
        # 初始化状态为全0
        init_state = np.zeros(self.state_size, int)
        return init_state

    # 执行动作并更新环境状态
    def step(self, action):
        global TDMA_counter

        # 初始化TDMA节点的奖励
        tdma_reward = 0
        # 初始化DRL节点的奖励
        agent_reward = 0
        # 初始化sa节点的奖励
        sa_reward = 0
        # 初始化总奖励
        reward = 0
        # 初始化观察结果
        observation_ = 0

        # 判断TDMA结点的动作
        if self.delay_tag > 0:  # 当TDMA节点距离AP比DRL节点远时
            if self.TDMAaction[TDMA_counter - self.delay_tag] == 1:
                TDMA_action = 1
            else:
                TDMA_action = 0
        else:
            if self.TDMAaction[TDMA_counter - self.delay_tag-10] == 1:
                TDMA_action = 1
            else:
                TDMA_action = 0

        # 不是TDMA发送的时隙，判断SA节点的z的大小，如果z>0，则传输
        # 判断TDMA节点的发送状态，如果TDMA发送，dqn不发送，
        if(TDMA_action==1):
            # TDMA节点都发送
            # DQN节点 不发送，但是碰撞
            if action ==1:
                # 传输，但不成功
                reward=0
                agent_reward=0
                tdma_reward=0
                sa_reward=0
                observation_ = -1
            else:
                # 无传输，成功
                reward = 0
                agent_reward = 0
                tdma_reward = 1
                sa_reward = 0
                observation_ = 2

        else: #TDMA_action==0
            # 两个TDMA节点都不发送，所以要参考SA的发送概率大小决定是否要发送
            if self.SAaction>=0 and self.SAaction<0.5 : # sa节点不发送
                if action==1:
                    # 传输，成功
                    reward = 1

                    agent_reward = 1
                    tdma_reward = 0
                    sa_reward = 0
                    observation_ = 1  # 传输，成功
                else:
                    # 未传输，也不成功
                    reward = 0
                    agent_reward = 0
                    tdma_reward = 0
                    sa_reward = 0
                    observation_ = -2 # 未传输，也不成功
            else: # SA发送
                if action == 1:
                    # 传输，但是不成功
                    reward = 0
                    agent_reward = 0
                    tdma_reward = 0
                    sa_reward = 0
                    observation_ = -1  # 传输，但未成功
                else:
                    # SA节点成功传输
                    reward = 1
                    agent_reward = 0
                    tdma_reward = 0
                    sa_reward = 1
                    observation_ = 2  # 未传输，但成功


        # 更新TDMA计数器
        TDMA_counter += 1
        # 如果计数器到达动作序列的末尾，则重置
        if TDMA_counter == len(self.TDMAaction):
            TDMA_counter = 0


        # 记录DRL节点的奖励
        self.agent_reward_list.append(agent_reward)
        # 记录TDMA节点的奖励
        self.tdma_reward_list.append(tdma_reward)
        # 记录sa节点的奖励
        self.sa_reward_list.append(sa_reward)
        # 记录总奖励
        self.reward_list.append(reward)
        # 记录DRL节点的观察结果
        self.DRL_observation_list.append(observation_)
        # 增加观察奖励计数器
        self.observation_reward_counter += 1
        # 如果观察奖励计数器达到2倍的DRL延迟，则返回最新的观察结果和奖励
        # 为什么要返回tdma结点的奖励？？？？？？？？
        if self.observation_reward_counter >= 2 * self.DRL_delay:
            counter11 = self.observation_reward_counter - (2 * self.DRL_delay)
            return self.DRL_observation_list[counter11], self.reward_list[counter11], self.agent_reward_list[counter11], self.tdma_reward_list[counter11],self.sa_reward_list[counter11]
        else:
            # 否则返回0
            return 0, 0, 0, 0, 0